#!/usr/bin/env python
#-*- coding: utf-8 -*-


'''
描述
1, 
查询该object表的分布，
可以查到当前该object表都分布在哪些个site上，rack上，node上，disk上。
比如我们记做mmap,
如果mmap中disk数目小于该object表应有的副本数，则退出, 执行完成。

2，
查询当前系统的diskmap信息，site, rack, node, disk。
我们记做map

3,
找出该objcet表还未分布的new_sites, 得到的sites可以为空列表。

4, 
给未分布到的sites分配副本。
从该object表已分布的site中查询，如果哪个site分布了超过一个副本，
就从该site中移动一个副本到新的site中。
用随机的方式从新的site中选择一个disk。
更新mmap, 删除被移动的disk, 添加新的disk。

5, 
根据最新的map和mmap,
找出该objcet表还未分布的new_racks, 得到的racks可以为空列表, 
并且该列表中的项为 (site, rack) 

6，
从该object表已分布的rack中查询， 如果某个rack1分布超过了一个副本, 
就把该rack中的一个副本移到到新的rack2中。

rack2 是从第5步中查询到的racks中选出来的。
从racks中选择rack时，如果rack的site和rack1的site不同，则优先级高,
site相同则优先级低。
在相同的优先级中随机选择一个rack。
从选择出来的rack中，随机选择一个disk。
更新mmap, 删除被移动的disk, 添加新的disk。

7,
根据最新的map和mmap,
找出该objcet表还未分布的new_nodes, 得到的nodes可以为空列表, 
并且该列表中的项为 (site, rack, node) 

8，
从该object表已分布的node中查询， 如果某个node1分布超过了一个副本, 
就把该node中的一个副本移到到新的node2中。

node2 是从第7步中查询到的nodes中选出来的。

从nodes中选择node时，
如果node的site和node1的site不同，则优先级高,site相同则优先级低。

如果node的(site, rack)和node1的(site, rack)不同，则优先级高,
(site, rack)相同则优先级低。

在相同的优先级中随机选择一个disk。
更新mmap, 删除被移动的disk, 添加新的disk。

9,
得到新的目标磁盘，需要把loc_master，放在第一位

9,
执行move操作
'''

import re
import random
import json
import errno

from utils import Exp
from diskmap import DiskMap, name2map, map2name
from diskmap_utils import cluster_init, cluster_add_disk, cluster_del_disk, \
    disks_of_cluster, disks_of_site, disks_of_rack, disks_of_node, \
    site_of_disks_max, rack_of_disks_max, node_of_disks_max, \
    get_random_rack, get_random_node, get_random_disk, \
    get_sites_of_gt_one_disk, get_racks_of_gt_one_disk, get_nodes_of_gt_one_disk


def racks_sort(cluster, racks):
    #降序 desc
    #racks = [[site, rack], ...]
    racks = sorted(
            racks, 
            cmp=lambda x,y:cmp(
                len(disks_of_rack(cluster, x[0], x[1])), 
                len(disks_of_rack(cluster, y[0], y[1])) 
                ), 
            reverse=True)

    return racks

def nodes_sort(cluster, nodes):
    #降序 desc
    #racks = [[site, rack, node], ...]
    nodes = sorted(
            nodes, 
            cmp=lambda x,y:cmp(
                len(disks_of_node(cluster, x[0], x[1], x[2])), 
                len(disks_of_node(cluster, y[0], y[1], y[2])) 
                ), 
            reverse=True)

    return nodes

def get_disk_with_site_priority(cluster, racks_new, site_lower):
    #racks_new = [(site, rack), ...]
    #return [site, rack, node, disk]
    racks_lower = []
    racks_high = []

    for x in racks_new:
        if x[0] == site_lower:
            racks_lower.append(x)
        else:
            racks_high.append(x)

    for x in [racks_high, racks_lower]:
        if x:
            i = random.randint(0, len(x)-1)
            rack = x[i]

            node = get_random_node(cluster, rack[0], rack[1])
            disk = get_random_disk(cluster, node[0], node[1], node[2])
            return disk

def get_disk_with_site_rack_priority(cluster, nodes_new, site_lower, rack_lower):
    #nodes_new = [(site, rack, node), ...]
    #return [site, rack, node, disk]
    nodes_lower = []
    nodes_high = []

    nodes_out_site = []
    for x in nodes_new:
        if x[0] == site_lower:
            continue
        nodes_out_site.append(x)

    nodes_out_rack = []
    for x in nodes_new:
        if x[0] == site_lower and x[1] == rack_lower:
            continue
        nodes_out_rack.append(x)

    nodes_in_rack = []
    for x in nodes_new:
        if x[0] == site_lower and x[1] == rack_lower:
            nodes_in_rack.append(x)

    for x in [nodes_out_site, nodes_out_rack, nodes_in_rack]:
        if x:
            i = random.randint(0, len(x)-1)
            node = x[i]

            disk = get_random_disk(cluster, node[0], node[1], node[2])
            return disk

def get_disk_skip_master_in_site(locs, site, loc_master):
    #return [site, rack, node, disk]
    rack = rack_of_disks_max(locs, site)
    node = node_of_disks_max(locs, rack[0], rack[1])
    disks = disks_of_node(locs, node[0], node[1], node[2])
    disk = None

    if len(disks) > 1:
        if map2name(disks[0]) == loc_master:
            disk = disks[1]
        else:
            disk = disks[0]
    else:
        if map2name(disks[0]) == loc_master:
            disks = disks_of_rack(locs, rack[0], rack[1])
            for x in disks:
                if not map2name(x) == loc_master:
                    disk = disks[1]
                    break

            if disk is None:
                disks = disks_of_site(locs, site)
                for x in disks:
                    if not map2name(x) == loc_master:
                        disk = disks[1]
                        break

        else:
            disk = disks[0]

    assert(disk is not None)
    return disk


def get_disk_skip_master_in_rack(locs, site, rack, loc_master):
    #return [site, rack, node, disk]
    node = node_of_disks_max(locs, site, rack)
    disks = disks_of_node(locs, node[0], node[1], node[2])
    disk = None

    if len(disks) > 1:
        if map2name(disks[0]) == loc_master:
            disk = disks[1]
        else:
            disk = disks[0]
    else:
        if map2name(disks[0]) == loc_master:
            disks = disks_of_rack(locs, site, rack)
            for x in disks:
                if not map2name(x) == loc_master:
                    disk = disks[1]
                    break

        else:
            disk = disks[0]

    assert(disk is not None)
    return disk

def get_disk_skip_master_in_node(locs, site, rack, node, loc_master):
    #return [site, rack, node, disk]
    disks = disks_of_node(locs, site, rack, node)
    disk = None

    if map2name(disks[0]) == loc_master:
        disk = disks[1]
    else:
        disk = disks[0]

    assert(disk is not None)
    return disk


def get_new_locs(lst, locs, loc_master):
    #在locs 长度假定就是replica, 
    #如果不是就要等待recover，然后在balance
    #return [map2name(x) for x in disks_of_cluster(locs)]
    _locs = {}

    for i in locs: 
        _locs.update({i:'normal'})

    locs = cluster_init(_locs)
    cluster = cluster_init(lst)

    locs_offline = list(set([map2name(x) for x in disks_of_cluster(locs)])
            - set([map2name(x) for x in disks_of_cluster(cluster)]))

    if locs_offline:
        msg = ("nodes_offline: %s" % (str(locs_offline)))
        raise Exp(errno.EINVAL, msg)

    sites_cluster = [x.name for x  in cluster.site.child]
    sites_locs = [x.name for x in locs.site.child]
    sites_new = list(set(sites_cluster) - set(sites_locs))

    for x in sites_new:
        sites_more = get_sites_of_gt_one_disk(locs)
        if not sites_more:
            break
        disk = get_disk_skip_master_in_site(locs, sites_more[0], loc_master)
        cluster_del_disk(locs, disk[0], disk[1], disk[2], disk[3])

        rack = get_random_rack(cluster, x)
        node = get_random_node(cluster, rack[0], rack[1])
        disk = get_random_disk(cluster, node[0], node[1], node[2])
        cluster_add_disk(locs, disk[0], disk[1], disk[2], disk[3])

    while True:
        racks_locs = []
        racks_cluster = []

        for site_ent in locs.site.child:
            for rack_ent in site_ent.child:
                racks_locs.append((site_ent.name, rack_ent.name))

        for site_ent in cluster.site.child:
            for rack_ent in site_ent.child:
                racks_cluster.append((site_ent.name, rack_ent.name))

        racks_new = list(set(racks_cluster) - set(racks_locs))

        if not racks_new:
            break

        racks_more = []
        for site_ent in locs.site.child:
            racks = get_racks_of_gt_one_disk(locs, site_ent.name)
            racks_more += racks

        if not racks_more:
            break

        racks_more = racks_sort(cluster, racks_more)

        x = racks_more[0]
        disk = get_disk_skip_master_in_rack(locs, x[0], x[1], loc_master)
        cluster_del_disk(locs, disk[0], disk[1], disk[2], disk[3])

        disk = get_disk_with_site_priority(cluster, racks_new, x[0])
        cluster_add_disk(locs, disk[0], disk[1], disk[2], disk[3])

    while True:
        nodes_locs = []
        nodes_cluster = []

        for site_ent in locs.site.child:
            for rack_ent in site_ent.child:
                for node_ent in rack_ent.child:
                    nodes_locs.append((site_ent.name, rack_ent.name, node_ent.name))

        for site_ent in cluster.site.child:
            for rack_ent in site_ent.child:
                for node_ent in rack_ent.child:
                    nodes_cluster.append((site_ent.name, rack_ent.name, node_ent.name))

        nodes_new = list(set(nodes_cluster) - set(nodes_locs))

        if not nodes_new:
            break

        nodes_more = []
        for site_ent in locs.site.child:
            for rack_ent in site_ent.child:
                nodes = get_nodes_of_gt_one_disk(locs, site_ent.name, rack_ent.name)
                nodes_more += nodes

        if not nodes_more:
            break

        nodes_more = nodes_sort(cluster, nodes_more)

        x = nodes_more[0]
        disk = get_disk_skip_master_in_node(locs, x[0], x[1], x[2], loc_master)
        cluster_del_disk(locs, disk[0], disk[1], disk[2], disk[3])

        disk = get_disk_with_site_rack_priority(cluster, nodes_new, x[0], x[1])
        cluster_add_disk(locs, disk[0], disk[1], disk[2], disk[3])

    new_disks = [map2name(x) for x in disks_of_cluster(locs)]

    rs = [loc_master]
    for i in new_disks:
        if i != loc_master:
            rs.append(i)

    assert(rs[0] == loc_master)
    return rs

def __test_add_node():
    lst = {
        'd1.r1.h1/2': 'normal', 
        'd1.r1.h1/0': 'normal', 
        'd1.r1.h1/1': 'normal', 
    
        'd2.r2.h2/0': 'meta', 
        'd2.r2.h2/1': 'normal', 
        'd2.r2.h2/2': 'normal', 
    
        'd3.r3.h3/0': 'meta', 
        'd3.r3.h3/1': 'normal', 
        'd3.r3.h3/2': 'normal',

        'd3.r3.h4/0': 'meta', 
        'd3.r3.h4/1': 'normal', 
        'd3.r3.h4/2': 'normal',
       }

    print("add node:")
    locs =  ['d1.r1.h1/0', 'd1.r1.h1/2', 'd2.r2.h2/1', 'd3.r3.h3/2']
    master = 'd1.r1.h1/2'
    locs_new = get_new_locs(lst, locs, master)

    assert(locs_new[0] == master)
    print("locs: ", sorted(locs)) 
    print("locs_new: ", sorted(locs_new)) 

def __test_add_rack():
    lst = {
        'd1.r1.h1/2': 'normal', 
        'd1.r1.h1/0': 'normal', 
        'd1.r1.h1/1': 'normal', 
    
        'd2.r2.h2/0': 'meta', 
        'd2.r2.h2/1': 'normal', 
        'd2.r2.h2/2': 'normal', 
    
        'd3.r3.h3/0': 'meta', 
        'd3.r3.h3/1': 'normal', 
        'd3.r3.h3/2': 'normal',

        'd3.r4.h4/0': 'meta', 
        'd3.r4.h4/1': 'normal', 
        'd3.r4.h4/2': 'normal',
       }

    print("add rack:")
    locs =  ['d1.r1.h1/0', 'd1.r1.h1/2', 'd2.r2.h2/1', 'd3.r3.h3/2']
    master = 'd1.r1.h1/2'
    locs_new = get_new_locs(lst, locs, master)

    assert(locs_new[0] == master)
    print("locs: ", sorted(locs)) 
    print("locs_new: ", sorted(locs_new)) 

def __test_add_site():
    lst = {
        'd1.r1.h1/2': 'normal', 
        'd1.r1.h1/0': 'normal', 
        'd1.r1.h1/1': 'normal', 
    
        'd2.r2.h2/0': 'meta', 
        'd2.r2.h2/1': 'normal', 
        'd2.r2.h2/2': 'normal', 
    
        'd3.r3.h3/0': 'meta', 
        'd3.r3.h3/1': 'normal', 
        'd3.r3.h3/2': 'normal',

        'd4.r4.h4/0': 'meta', 
        'd4.r4.h4/1': 'normal', 
        'd4.r4.h4/2': 'normal',
       }

    print("add site:")
    locs =  ['d1.r1.h1/0', 'd1.r1.h1/2', 'd2.r2.h2/1', 'd3.r3.h3/2']
    master = 'd1.r1.h1/2'
    locs_new = get_new_locs(lst, locs, master)

    assert(locs_new[0] == master)
    print("locs: ", sorted(locs)) 
    print("locs_new: ", sorted(locs_new)) 

def __test_add_more_node():
    lst = {
        'd1.r1.h1/2': 'normal', 
        'd1.r1.h1/0': 'normal', 
        'd1.r1.h1/1': 'normal', 
    
        'd2.r2.h2/0': 'meta', 
        'd2.r2.h2/1': 'normal', 
        'd2.r2.h2/2': 'normal', 
    
        'd3.r3.h3/0': 'meta', 
        'd3.r3.h3/1': 'normal', 
        'd3.r3.h3/2': 'normal',

        'd3.r3.h4/0': 'meta', 
        'd3.r3.h4/1': 'normal', 
        'd3.r3.h4/2': 'normal',

        'd3.r3.h5/0': 'meta', 
        'd3.r3.h5/1': 'normal', 
        'd3.r3.h5/2': 'normal',

        'd3.r3.h6/0': 'meta', 
        'd3.r3.h6/1': 'normal', 
        'd3.r3.h6/2': 'normal',
       }

    print("add more node:")
    locs =  ['d1.r1.h1/0', 'd1.r1.h1/2', 'd2.r2.h2/1', 'd3.r3.h3/2']
    master = 'd1.r1.h1/2'
    locs_new = get_new_locs(lst, locs, master)

    assert(locs_new[0] == master)
    print("locs: ", sorted(locs)) 
    print("locs_new: ", sorted(locs_new)) 



def __test_add_more_rack():
    lst = {
        'd1.r1.h1/2': 'normal', 
        'd1.r1.h1/0': 'normal', 
        'd1.r1.h1/1': 'normal', 
    
        'd2.r2.h2/0': 'meta', 
        'd2.r2.h2/1': 'normal', 
        'd2.r2.h2/2': 'normal', 
    
        'd3.r3.h3/0': 'meta', 
        'd3.r3.h3/1': 'normal', 
        'd3.r3.h3/2': 'normal',

        'd3.r4.h4/0': 'meta', 
        'd3.r4.h4/1': 'normal', 
        'd3.r4.h4/2': 'normal',

        'd3.r5.h4/0': 'meta', 
        'd3.r5.h4/1': 'normal', 
        'd3.r5.h4/2': 'normal',

        'd3.r6.h4/0': 'meta', 
        'd3.r6.h4/1': 'normal', 
        'd3.r6.h4/2': 'normal',
       }

    print("add more rack:")
    locs =  ['d1.r1.h1/0', 'd1.r1.h1/2', 'd2.r2.h2/1', 'd3.r3.h3/2']
    master = 'd1.r1.h1/2'
    locs_new = get_new_locs(lst, locs, master)

    assert(locs_new[0] == master)
    print("locs: ", sorted(locs)) 
    print("locs_new: ", sorted(locs_new)) 

def __test_add_more_site():
    lst = {
        'd1.r1.h1/2': 'normal', 
        'd1.r1.h1/0': 'normal', 
        'd1.r1.h1/1': 'normal', 
    
        'd2.r2.h2/0': 'meta', 
        'd2.r2.h2/1': 'normal', 
        'd2.r2.h2/2': 'normal', 
    
        'd3.r3.h3/0': 'meta', 
        'd3.r3.h3/1': 'normal', 
        'd3.r3.h3/2': 'normal',

        'd4.r4.h4/0': 'meta', 
        'd4.r4.h4/1': 'normal', 
        'd4.r4.h4/2': 'normal',

        'd5.r5.h4/0': 'meta', 
        'd5.r5.h4/1': 'normal', 
        'd5.r5.h4/2': 'normal',

        'd6.r6.h4/0': 'meta', 
        'd6.r6.h4/1': 'normal', 
        'd6.r6.h4/2': 'normal',
       }

    print("add more site:")
    locs =  ['d1.r1.h1/0', 'd1.r1.h1/2', 'd2.r2.h2/1', 'd3.r3.h3/2']
    master = 'd1.r1.h1/2'
    locs_new = get_new_locs(lst, locs, master)

    assert(locs_new[0] == master)
    print("locs: ", sorted(locs)) 
    print("locs_new: ", sorted(locs_new)) 

def test():
    __test_add_node()
    __test_add_rack()
    __test_add_site()
    __test_add_more_node()
    __test_add_more_rack()
    __test_add_more_site()

if __name__ == "__main__":
    test()
